import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
from supervised_vgg16 import super_vgg16
from supervised_vgg16 import prepare_data,data_preprocessing
from foolbox.models import TensorFlowModel
from foolbox.criteria import Misclassification
from foolbox.attacks import LBFGSAttack

train_x, train_y, test_x, test_y = prepare_data()
train_x, test_x = data_preprocessing(train_x, test_x)

image = tf.placeholder(tf.float32,shape=(None,32,32,3),name='input_x')
y_ = tf.placeholder(tf.float32,shape=[None,10],name='input_y')
# train_flag =  tf.placeholder(tf.bool,name='fll')
# keep_prob = tf.placeholder(tf.float32,name='prob')
loss,logits = super_vgg16(image,y_,keep_prob=1.0,phase=False)
saver = tf.train.Saver()

with tf.Session() as session:
    saver.restore(session, tf.train.latest_checkpoint('./supervised_vgg16_model'))
    print('finish loading model!')

    model = TensorFlowModel(image, logits, bounds=(-255, 255))
    criterion = Misclassification()
    attack = LBFGSAttack(model,criterion)

    # images = test_x[1]
    # label = np.argmax(model.forward_one(images))
    # True_label = np.argmax(test_y[1])
    # print(label)
    # print(True_label)
    # adversarial = attack(images,label)
    # print(adversarial.shape)
    # print(model.forward_one(adversarial))
    per_dataset = np.zeros((1000,32,32,3))
    for i in range(1000):
        images = test_x[i]
        label = np.argmax(model.predictions(images))
        adversarial = attack(images,label)
        per_dataset[i] = adversarial
        print("makeing picture" + str(i))
    np.save("LBFGS_vgg16_all",per_dataset)
    # plt.figure()
    # plt.subplot(1, 3, 1)
    # plt.imshow(images)
    #
    # plt.subplot(1, 3, 2)
    # plt.imshow(adversarial)
    #
    # plt.subplot(1, 3, 3)
    # plt.imshow(adversarial - images)
    # plt.show()